package main

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
)

func main() {
	if err := filepath.WalkDir(".", walk); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("✅ All notebooks cleaned.")
}

func walk(path string, d fs.DirEntry, err error) error {
	if err != nil {
		return err
	}
	if filepath.Ext(path) != ".ipynb" {
		return nil
	}
	return cleanNotebook(path)
}

func cleanNotebook(path string) error {
	raw, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("failed to read %s: %w", path, err)
	}

	var root map[string]interface{}
	if err := json.Unmarshal(raw, &root); err != nil {
		return fmt.Errorf("failed to parse %s: %w", path, err)
	}

	cells, ok := root["cells"].([]interface{})
	if !ok { // 没有 cells 字段
		return nil
	}

	modified := false
	for _, c := range cells {
		cell, ok := c.(map[string]interface{})
		if !ok || cell["cell_type"] != "code" {
			continue
		}
		// 只清这两个键
		if len(cell["outputs"].([]interface{})) > 0 || cell["execution_count"] != nil {
			cell["outputs"] = []interface{}{}
			cell["execution_count"] = nil
			modified = true
		}
	}

	if !modified {
		return nil
	}

	out, err := json.MarshalIndent(&root, "", " ")
	if err != nil {
		return fmt.Errorf("failed to marshal %s: %w", path, err)
	}

	if err := os.WriteFile(path, out, 0644); err != nil {
		return fmt.Errorf("failed to write %s: %w", path, err)
	}

	fmt.Printf("🧹 Cleaned: %s\n", path)
	return nil
}

/* 编译
go build -ldflags="-s -w" -trimpath -o nbstripdir.exe nbstripdir.go
./clean.exe        # 递归当前目录
*/
